import cv2
import os
import numpy as np
import matplotlib.pyplot as plt


def count_val(ori, filename):
    count = {}
    for i in range(256):
        count[i] = 0

    # 统计各灰度值的频率
    for row in ori:
        for col in row:
            count[col[0]] += 1

    # 绘制直方图并保存
    plt.bar(count.keys(), count.values())
    plt.savefig(os.path.abspath('./res/exp1/' + filename + '.png'), format='png')
    plt.show()
    return count


def process(ori, filename):
    # 读取图片
    ori_width = len(ori[0])
    ori_height = len(ori)
    count = count_val(ori, 'ori' + filename)

    # 进行频率累加
    acc = {0: count[0]}
    for i in range(1, 256):
        acc[i] = acc[i - 1] + count[i]

    # 计算对应关系
    correspond = {}
    m = ori_height * ori_width
    for key, val in acc.items():
        correspond[key] = int(val / m * 255 + 0.5)

    # 通过对应关系修改相关像素值
    new = np.zeros_like(ori)
    for i in range(ori_height):
        for j in range(ori_width):
            new[i, j, :] = correspond[ori[i, j, 0]]

    # 统计均衡化后的频率，同时输出直方图
    count_val(new, 'new' + filename)
    cv2.imwrite('res/exp1/res_' + filename + '.jpg', new)
    return new


if __name__ == '__main__':
    """
    文件皆存放在./res/exp1文件夹下，命名规则如下：
        以new开头的为均衡化之后的图像直方图
        以ori开头的为均衡化之前的图像直方图
        以res开头的为均衡化之后的图像
    """
    ori1 = cv2.imread('./test/test1_1.jpg')
    process(ori1, 'test1_1')

    ori2 = cv2.imread('./test/test1_2.jpg')
    width = len(ori2[0])
    height = len(ori2)
    for i in range(height):
        for j in range(width):
            ori2[i, j, :] = np.sum(ori2[i, j, :]) / 3

    cv2.imwrite('./res/exp1/gray.jpg', ori2)

    ori3 = cv2.imread('./test/test1_3.jpg')
    process(ori3, 'test1_3')
